#include "tree_node.h"
#include <cassert>
#include <stack>

std::shared_ptr<Node> min(const std::shared_ptr<Node>& root){
	std::shared_ptr<Node> curr = root;
	while(curr->left)
		curr = curr->left;
	return curr;
}

std::shared_ptr<Node> max(const std::shared_ptr<Node>& root){
	std::shared_ptr<Node> curr = root;
	while(curr->right)
		curr = curr->right;
	return curr;
}

std::shared_ptr<Node> predecessor(const std::shared_ptr<Node>& node){
	if(node->left)
		return max(node->left);
	std::shared_ptr<Node> y = node->parent;
	std::shared_ptr<Node> x = node;
	while (y != nullptr){
		if(y->right == x)
			return y;
		x = y;
		y = y->parent;
	}
	return nullptr;
}
//assumption: node != nullptr
std::shared_ptr<Node> successor(const std::shared_ptr<Node>& node){
	if(node->right)
		return min(node->right);
	std::shared_ptr<Node> y = node->parent;
	std::shared_ptr<Node> x = node;
	while(y != nullptr) {
		if(y->left == x)
			return y;
		x = y;
		y = y->parent;
	}
	return nullptr;
}
//return empty ptr if kth_min doesn't exist
std::shared_ptr<Node> kth_min(const std::shared_ptr<Node>& root, int k){
	std::shared_ptr<Node> curr = min(root);
	for ( int i = 1; i <= k-1; ++i){
		std::shared_ptr<Node> next = successor(curr);
		curr = next;
		if(!curr)
			break;
	}
	return curr;
}

std::shared_ptr<Node> kth_min_2(const std::shared_ptr<Node>& root, int k) {
	std::shared_ptr<Node> curr = root;
	std::stack<std::shared_ptr<Node> > s;
	while(curr != nullptr || !s.empty()){
		if(curr!= nullptr) {
			s.push(curr);
			curr = curr->left;
		}
		else {
			curr = s.top();
			s.pop();
			if(k-- == 1)
				return curr;
			curr = curr->right;
		}
	}
	return nullptr;
}

std::shared_ptr<Node> kth_max(const std::shared_ptr<Node>& root, int k){
	std::shared_ptr<Node> curr = max(root);
	for (int i = 1; i <= k-1; ++i) {
		curr = predecessor(curr);
		if(!curr)
			break;
	}
	return curr;
}

std::shared_ptr<Node> find(const std::shared_ptr<Node>& root, int key){
	std::shared_ptr<Node> curr = root;
	while(curr != nullptr) {
		if( curr->key==key )
			return curr;
		else if (curr->key > key) 
			curr = curr->left;
		else
			curr = curr->right;
	}
	return nullptr;
}

void test_min_max(){
	std::string str = "3 1 0 # # 2 # # 5 4 # # 6 # #";
	std::shared_ptr<Node> root = preorder_deserialize(str);
	print_tree(root);
	std::cout << "The minimum nodes is " << min(root)->key << std::endl;
	std::cout << "The maximum nodes is " << max(root)->key << std::endl;

	int key;
	std::cout << "Enter the key:\n";
	while(std::cin >> key){
		std::shared_ptr<Node> node = find(root,key);
		if(!node){
			std::cout << "Can't find the key\n";
			return;
		}
		std::shared_ptr<Node> pred = predecessor(node);
		if(pred)
			std::cout << "The predecessor is " << pred->key << std::endl;
		else
			std::cout << "No predecessor\n";
		std::shared_ptr<Node> succ = successor(node);
		if(succ) 
			std::cout << "The successor is " << succ->key << std::endl;
		else
			std::cout << "No successor\n";
	}
}

void test_kth_min_max(){
	std::string str = "3 1 0 # # 2 # # 5 4 # # 6 # #";
	std::shared_ptr<Node> root = preorder_deserialize(str);
	print_tree(root);
	std::cout << "Enter the kth:\n";
	int k;
	while(std::cin >> k){
		assert(k>0);
		std::shared_ptr<Node> kth = kth_min(root,k);
		if(kth)
			std::cout << "the kth minimum is " << kth->key << std::endl;
		else
			std::cout << "There is no " << k << "th smallest node\n";

		kth = kth_min_2(root,k);
		if(kth)
			std::cout << "the kth minimum is " << kth->key << std::endl;
		else
			std::cout << "There is no " << k << "th smallest node\n";

		kth = kth_max(root,k);
		if(kth)
			std::cout << "the kth maximum is " << kth->key << std::endl;
		else
			std::cout << "There is no " << k << "th smallest node\n";
	}
}